# Venus Data Processing Batch Script with Scale Parameter
# Saved as .ps1 file (e.g., process_venus.ps1) and run with administrator privileges

# Parse command line arguments
param(
    [string]$scale = "tiny"
)

# Set UTF-8 encoding to prevent Chinese character display issues
$OutputEncoding = [System.Text.Encoding]::UTF8
[Console]::OutputEncoding = [System.Text.Encoding]::UTF8

# Define scale to number mapping
$scaleMap = @{
    "tiny" = 5
    "small" = 20
    "medium" = 80
    "large" = 320
    "huge" = 1280
    "full" = 4600
}

# Validate scale parameter
if (-not $scaleMap.ContainsKey($scale)) {
    Write-Host "Invalid scale parameter! Available options: tiny, small, medium, large, huge, full" -ForegroundColor Red
    exit 1
}

# Get the number of videos to process based on scale
$numVideos = $scaleMap[$scale]
Write-Host "Selected scale: $scale, processing $numVideos videos" -ForegroundColor Cyan

# Define directory names based on scale
$outputDir = "data/venus$numVideos"
$targetDir = "data/venus_$scale"

Write-Host "Starting Venus data processing pipeline for scale: $scale..." -ForegroundColor Green

# A: Split and resize videos
Write-Host "A1. Splitting and resizing videos..." -ForegroundColor Yellow
python data/src/preprocess_videos.py --input E:/video_small --output $outputDir --shortside 224 --num $numVideos --gpu 2
if ($LASTEXITCODE -ne 0) {
    Write-Host "Video preprocessing failed!" -ForegroundColor Red
    exit 1
}

Write-Host "A2. Copying processed videos..." -ForegroundColor Yellow
Copy-Item -Path $outputDir -Destination $targetDir -Recurse -Force
if (-not (Test-Path $targetDir)) {
    Write-Host "Video copy failed!" -ForegroundColor Red
    exit 1
}

# B: Generate action and speed annotation files
Write-Host "B. Generating action and speed annotation files..." -ForegroundColor Yellow

Write-Host "B1. Generating action annotations..." -ForegroundColor Yellow
python data/src/generate_annotations.py --video_dir $targetDir --output "${targetDir}/venus_$scale.txt" --type category
if ($LASTEXITCODE -ne 0) {
    Write-Host "Action annotation generation failed!" -ForegroundColor Red
    exit 1
}

Write-Host "B2. Generating speed annotations..." -ForegroundColor Yellow
python data/src/generate_annotations.py --video_dir $targetDir --output "${targetDir}/venus_$scale.txt" --type speed
if ($LASTEXITCODE -ne 0) {
    Write-Host "Speed annotation generation failed!" -ForegroundColor Red
    exit 1
}

# C: Clean action annotation files and videos
Write-Host "C. Cleaning action annotation files and videos..." -ForegroundColor Yellow

Write-Host "C1. Cleaning training set..." -ForegroundColor Yellow
python data/src/dataset_cleaner.py --annotation_file "${targetDir}/venus_${scale}_train_cate.txt" --video_base_dir "${targetDir}/train" --seed 59
if ($LASTEXITCODE -ne 0) {
    Write-Host "Training set cleaning failed!" -ForegroundColor Red
    exit 1
}

Write-Host "C2. Cleaning validation set..." -ForegroundColor Yellow
python data/src/dataset_cleaner.py --annotation_file "${targetDir}/venus_${scale}_val_cate.txt" --video_base_dir "${targetDir}/val" --seed 59
if ($LASTEXITCODE -ne 0) {
    Write-Host "Validation set cleaning failed!" -ForegroundColor Red
    exit 1
}

# D: Synchronize speed annotation files
Write-Host "D. Synchronizing speed annotation files..." -ForegroundColor Yellow

Write-Host "D1. Synchronizing training set speed annotations..." -ForegroundColor Yellow
python data/src/sync_annotations.py --action_anno "${targetDir}/venus_${scale}_train_cate.txt" --speed_anno "${targetDir}/venus_${scale}_train_speed.txt"
if ($LASTEXITCODE -ne 0) {
    Write-Host "Training set speed annotation synchronization failed!" -ForegroundColor Red
    exit 1
}

Write-Host "D2. Synchronizing validation set speed annotations..." -ForegroundColor Yellow
python data/src/sync_annotations.py --action_anno "${targetDir}/venus_${scale}_val_cate.txt" --speed_anno "${targetDir}/venus_${scale}_val_speed.txt"
if ($LASTEXITCODE -ne 0) {
    Write-Host "Validation set speed annotation synchronization failed!" -ForegroundColor Red
    exit 1
}

Write-Host "All processing steps for scale '$scale' completed successfully!" -ForegroundColor Green